"""
Emittance analysis and visualization
Creates 3x3 grid plots for beam parameters vs frequency
"""

import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.signal import find_peaks
from utils.emittance_hanlder import process_files_for_emittance_analysis
from utils.console_hanlder import console
from utils.csv_hanlder import select_and_read_files


def create_emitance_vs_frequency_3x3_plots_linear(
    f, rms_emittance_x, rms_emittance_y, rms_gm_xy, geometric_x, geometric_y, geometric_gm_xy,
    normalized_x, normalized_y, normalized_gm_xy, area,
    figsize=(24, 18), fig_name=None, smooth_window=21, target_valley_count=(3, 5)):
    """
    Create a 3×3 grid of scatter plots for beam parameters vs frequency,
    automatically detecting Local Minima (local minima) and summarizing correlations.
    
    Features:
    - Smooths each parameter curve to reduce noise.
    - Detects 3–5 Local Minima per parameter using adaptive prominence threshold.
    - Marks valleys with red triangles (no legend for clean visuals).
    - Prints a detailed table of detected valleys and parameter statistics.
    - Displays a correlation heatmap (including convex hull area).
    """
    fig = plt.figure(figsize=figsize)
    fig.suptitle('Beam Parameters Analysis vs Frequency', fontsize=20, fontweight='bold')
    gs = fig.add_gridspec(3, 3, hspace=0.4, wspace=0.3)

    data_configs = [
        (rms_emittance_x, 'RMS Emittance X', gs[0, 0]),
        (rms_emittance_y, 'RMS Emittance Y', gs[0, 1]),
        (rms_gm_xy, 'RMS GM XY', gs[0, 2]),
        (geometric_x, 'Geometric X', gs[1, 0]),
        (geometric_y, 'Geometric Y', gs[1, 1]),
        (geometric_gm_xy, 'Geometric GM XY', gs[1, 2]),
        (normalized_x, 'Normalized X', gs[2, 0]),
        (normalized_y, 'Normalized Y', gs[2, 1]),
        (normalized_gm_xy, 'Normalized GM XY', gs[2, 2])
    ]

    stats = []

    for data, title, subplot_spec in data_configs:
        ax = fig.add_subplot(subplot_spec)
        
        # Sort and smooth
        df_temp = pd.DataFrame({'f': f, 'y': data}).sort_values('f')
        y_smooth = df_temp['y'].rolling(window=smooth_window, center=True, min_periods=1).mean()

        # Initial prominence guess
        prominence = 0.05 * (y_smooth.max() - y_smooth.min())
        valleys, _ = find_peaks(-y_smooth, prominence=prominence, distance=len(y_smooth)//40)

        # Adjust to target valley count
        # If too many valleys → increase prominence
        while len(valleys) > target_valley_count[1]:
            prominence *= 1.2
            valleys, _ = find_peaks(-y_smooth, prominence=prominence, distance=len(y_smooth)//40)
        # If too few valleys → decrease prominence
        while len(valleys) < target_valley_count[0]:
            prominence /= 1.2
            valleys, _ = find_peaks(-y_smooth, prominence=prominence, distance=len(y_smooth)//40)

        # Scatter plot with color mapping
        scatter = ax.scatter(df_temp['f'], df_temp['y'], c=df_temp['y'], cmap='plasma', s=30, alpha=0.5)
        plt.colorbar(scatter, ax=ax, shrink=0.8)

        # Mark detected valleys
        ax.scatter(df_temp['f'].iloc[valleys], y_smooth.iloc[valleys],
                   color='red', marker='v', s=100, zorder=5, edgecolors='white', linewidth=1.5, label='Local Minima')

        ax.set_title(f'{title} vs Frequency', fontsize=12, fontweight='bold')
        ax.set_xlabel('Frequency [Hz]')
        ax.set_ylabel(f'{title} (linear scale)')
        ax.legend()
        ax.grid(True, alpha=0.3)

        # Store detected points
        for v in valleys:
            stats.append((title, df_temp['f'].iloc[v], y_smooth.iloc[v]))

    plt.tight_layout()
    plt.subplots_adjust(top=0.94)

    if fig_name:
        plt.savefig(fig_name, dpi=300, bbox_inches='tight')
        console.log(f"Plot saved to: {fig_name}")

    plt.close(fig)

    # === Save detected valleys to file only ===
    console.log_detailed_results("=" * 100)
    console.log_detailed_results("检测到的局部最小值")
    console.log_detailed_results("=" * 100)
    console.log_detailed_results(f"{'参数':<20} {'频率':<15} {'数值':<12}")
    console.log_detailed_results("-" * 100)
    for p, freq, val in stats:
        console.log_detailed_results(f"{p:<20} {freq:<15.0f} {val:<12.2e}")

    # === Save correlation matrix to file only ===
    console.log_detailed_results("\n" + "=" * 100)
    console.log_detailed_results("相关系数矩阵 (含频率)")
    console.log_detailed_results("=" * 100)

    df_corr = pd.DataFrame({
        'frequency': f,
        'hull_area': area,
        'rms_emittance_x': rms_emittance_x,
        'rms_emittance_y': rms_emittance_y,
        'rms_gm_xy': rms_gm_xy,
        'geometric_x': geometric_x,
        'geometric_y': geometric_y,
        'geometric_gm_xy': geometric_gm_xy,
        'normalized_x': normalized_x,
        'normalized_y': normalized_y,
        'normalized_gm_xy': normalized_gm_xy
    })

    corr_matrix = df_corr.corr(method='pearson')

    plt.figure(figsize=(10, 8))
    sns.heatmap(corr_matrix, annot=True, fmt=".2f", cmap='coolwarm', square=True, cbar_kws={'shrink': .8})
    plt.title("Correlation Matrix of Beam Parameters", fontsize=14, fontweight='bold')
    plt.tight_layout()
    plt.savefig('result/beam_parameters_corr_matrix.png', dpi=300, bbox_inches='tight')
    console.log("相关系数矩阵已保存到: result/beam_parameters_corr_matrix.png")
    plt.close(fig)

    return corr_matrix


def create_emitance_vs_frequency_3x3_plots_log(
    f, rms_emittance_x, rms_emittance_y, rms_gm_xy, geometric_x, geometric_y, geometric_gm_xy,
    normalized_x, normalized_y, normalized_gm_xy, area,
    figsize=(24, 18), fig_name=None, smooth_window=21, target_valley_count=(3, 5)):
    """
    Create a 3×3 grid of scatter plots for beam parameters vs frequency using LOG scale,
    automatically detecting Local Minima (local minima) and summarizing correlations.
    
    Features:
    - Smooths each parameter curve to reduce noise.
    - Detects 3–5 Local Minima per parameter using adaptive prominence threshold.
    - Marks valleys with red triangles (no legend for clean visuals).
    - Prints a detailed table of detected valleys and parameter statistics.
    - Displays a correlation heatmap (including convex hull area).
    - Uses logarithmic scale for y-axis.
    """
    fig = plt.figure(figsize=figsize)
    gs = fig.add_gridspec(3, 3, hspace=0.4, wspace=0.3)

    data_configs = [
        (rms_emittance_x, 'RMS Emittance X [mm·rad]', gs[0, 0]),
        (rms_emittance_y, 'RMS Emittance Y [mm·rad]', gs[0, 1]),
        (rms_gm_xy, 'RMS GM XY [mm·rad]', gs[0, 2]),
        (geometric_x, 'Geometric X [mm·rad]', gs[1, 0]),
        (geometric_y, 'Geometric Y [mm·rad]', gs[1, 1]),
        (geometric_gm_xy, 'Geometric GM XY [mm·rad]', gs[1, 2]),
        (normalized_x, 'Normalized X [mm·rad]', gs[2, 0]),
        (normalized_y, 'Normalized Y [mm·rad]', gs[2, 1]),
        (normalized_gm_xy, 'Normalized GM XY [mm·rad]', gs[2, 2])
    ]

    stats = []

    for data, title, subplot_spec in data_configs:
        ax = fig.add_subplot(subplot_spec)
        
        # Sort and smooth
        df_temp = pd.DataFrame({'f': f, 'y': data}).sort_values('f')
        y_smooth = df_temp['y'].rolling(window=smooth_window, center=True, min_periods=1).mean()

        # Initial prominence guess
        prominence = 0.05 * (y_smooth.max() - y_smooth.min())
        valleys, _ = find_peaks(-y_smooth, prominence=prominence, distance=len(y_smooth)//40)

        # Adjust to target valley count
        # If too many valleys → increase prominence
        while len(valleys) > target_valley_count[1]:
            prominence *= 1.2
            valleys, _ = find_peaks(-y_smooth, prominence=prominence, distance=len(y_smooth)//40)
        # If too few valleys → decrease prominence
        while len(valleys) < target_valley_count[0]:
            prominence /= 1.2
            valleys, _ = find_peaks(-y_smooth, prominence=prominence, distance=len(y_smooth)//40)

        # Log scale plot with colorbar
        scatter = ax.scatter(df_temp['f'], df_temp['y'], c=df_temp['y'], cmap='plasma', s=30, alpha=0.7)
        plt.colorbar(scatter, ax=ax, shrink=0.8)

        # Mark detected valleys
        ax.scatter(df_temp['f'].iloc[valleys], y_smooth.iloc[valleys],
                   color='red', marker='v', s=100, zorder=5, edgecolors='white', linewidth=1.5, label='Local Minima')

        ax.set_title(f'{title} vs Frequency', fontsize=12, fontweight='bold')
        ax.set_xlabel('Frequency [Hz]')
        ax.set_ylabel(f'{title} (log scale)')
        ax.set_yscale('log')  # 设置对数尺度
        ax.legend()
        ax.grid(True, alpha=0.3)

        # Store detected points
        for v in valleys:
            stats.append((title, df_temp['f'].iloc[v], y_smooth.iloc[v]))

    plt.tight_layout()
    plt.subplots_adjust(top=0.94)

    if fig_name:
        plt.savefig(fig_name, dpi=300, bbox_inches='tight')
        console.log(f"Plot saved to: {fig_name}")

    plt.close(fig)

    # === Save detected valleys to file only ===
    console.log_detailed_results("=" * 100)
    console.log_detailed_results("检测到的局部最小值 (对数尺度)")
    console.log_detailed_results("=" * 100)
    console.log_detailed_results(f"{'参数':<20} {'频率':<15} {'数值':<12}")
    console.log_detailed_results("-" * 100)
    for p, freq, val in stats:
        console.log_detailed_results(f"{p:<20} {freq:<15.0f} {val:<12.2e}")

    # === Save correlation matrix to file only ===
    console.log_detailed_results("\n" + "=" * 100)
    console.log_detailed_results("相关系数矩阵 (含频率) - 对数尺度")
    console.log_detailed_results("=" * 100)

    df_corr = pd.DataFrame({
        'frequency': f,
        'hull_area': area,
        'rms_emittance_x': rms_emittance_x,
        'rms_emittance_y': rms_emittance_y,
        'rms_gm_xy': rms_gm_xy,
        'geometric_x': geometric_x,
        'geometric_y': geometric_y,
        'geometric_gm_xy': geometric_gm_xy,
        'normalized_x': normalized_x,
        'normalized_y': normalized_y,
        'normalized_gm_xy': normalized_gm_xy
    })

    corr_matrix = df_corr.corr(method='pearson')

    plt.figure(figsize=(10, 8))
    sns.heatmap(corr_matrix, annot=True, fmt=".2f", cmap='coolwarm', square=True, cbar_kws={'shrink': .8})
    plt.title("Correlation Matrix of Beam Parameters (Log Scale)", fontsize=14, fontweight='bold')
    plt.tight_layout()
    plt.savefig('result/beam_parameters_corr_matrix_logscale.png', dpi=300, bbox_inches='tight')
    console.log("相关系数矩阵已保存到: result/beam_parameters_corr_matrix_logscale.png")
    plt.close(fig)

    return corr_matrix


def plot_emittance_analysis_linear(dfs, filenames=None, save_path='result/emittance_vs_frequency_lineascale.png'):
    """
    Main function to plot emittance analysis using LINEAR scale
    
    Args:
        dfs (list): List of DataFrames
        filenames (list): List of filenames (optional)
        save_path (str): Path to save the plot
    """
    console.log_section("发射度分析 (线性尺度)")
    
    try:
        # Process files to get parameters
        data = process_files_for_emittance_analysis(dfs, filenames)
        
        # Create 3x3 plots (linear scale)
        corr_matrix = create_emitance_vs_frequency_3x3_plots_linear(
            f=data['frequencies'],
            rms_emittance_x=data['rms_emittance_x'],
            rms_emittance_y=data['rms_emittance_y'],
            rms_gm_xy=data['rms_gm_xy'],
            geometric_x=data['geometric_x'],
            geometric_y=data['geometric_y'],
            geometric_gm_xy=data['geometric_gm_xy'],
            normalized_x=data['normalized_x'],
            normalized_y=data['normalized_y'],
            normalized_gm_xy=data['normalized_gm_xy'],
            area=data['area'],
            fig_name=save_path
        )
        
        console.log(f"发射度分析 (线性尺度) 完成。图表已保存到: {save_path}")
        
        return corr_matrix
        
    except Exception as e:
        console.log(f"发射度分析 (线性尺度) 出错: {str(e)}")
        return None


def plot_emittance_analysis_log(dfs, filenames=None, save_path='result/emittance_vs_frequency_logcale.png'):
    """
    Main function to plot emittance analysis using LOG scale
    
    Args:
        dfs (list): List of DataFrames
        filenames (list): List of filenames (optional)
        save_path (str): Path to save the plot
    """
    console.log_section("发射度分析 (对数尺度)")
    
    try:
        # Process files to get parameters
        data = process_files_for_emittance_analysis(dfs, filenames)
        
        # Create 3x3 plots (log scale)
        corr_matrix = create_emitance_vs_frequency_3x3_plots_log(
            f=data['frequencies'],
            rms_emittance_x=data['rms_emittance_x'],
            rms_emittance_y=data['rms_emittance_y'],
            rms_gm_xy=data['rms_gm_xy'],
            geometric_x=data['geometric_x'],
            geometric_y=data['geometric_y'],
            geometric_gm_xy=data['geometric_gm_xy'],
            normalized_x=data['normalized_x'],
            normalized_y=data['normalized_y'],
            normalized_gm_xy=data['normalized_gm_xy'],
            area=data['area'],
            fig_name=save_path
        )
        
        console.log(f"发射度分析 (对数尺度) 完成。图表已保存到: {save_path}")
        
        return corr_matrix
        
    except Exception as e:
        console.log(f"发射度分析 (对数尺度) 出错: {str(e)}")
        return None


def plot_emittance_analysis(dfs, filenames=None, save_path='result/emittance_analysis.png'):
    """
    Main function to plot emittance analysis (default: linear scale)
    
    Args:
        dfs (list): List of DataFrames
        filenames (list): List of filenames (optional)
        save_path (str): Path to save the plot
    """
    return plot_emittance_analysis_linear(dfs, filenames, save_path)


def main():
    """Main entry point for emittance analysis"""
    console.log("发射度分析工具")
    console.log("=" * 50)
    
    # Select files
    dfs, filenames = select_and_read_files(file_type='csv')
    
    if len(dfs) == 0:
        console.log("未找到有效文件。退出分析。")
        return
    
    console.log(f"处理 {len(dfs)} 个文件进行发射度分析")
    
    # Run analysis
    plot_emittance_analysis(dfs, filenames)


if __name__ == "__main__":
    main()
